package com.kanq.controller;

import com.kanq.common.utils.HanlpUtils;
import com.kanq.common.utils.SseUtil;
import com.kanq.pojo.entity.KanqAiFile;
import com.kanq.pojo.entity.KanqAiFileChunk;
import com.kanq.pojo.request.ChatRequest;
import com.kanq.pojo.response.UploadFileResponse;
import com.kanq.service.KanqAiFileChunkService;
import com.kanq.service.KanqAiFileService;
import com.kanq.service.LLMService;
import com.kanq.service.QdrantService;
import com.kanq.tool.R;
import io.qdrant.client.grpc.Points;
import jakarta.annotation.Resource;
import lombok.extern.slf4j.Slf4j;
import org.apache.ibatis.annotations.Param;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.web.bind.annotation.*;
import org.springframework.web.multipart.MultipartFile;
import org.springframework.web.servlet.mvc.method.annotation.SseEmitter;

import java.util.List;
import java.util.UUID;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.Executors;

@Slf4j
@RestController
@RequestMapping("/llm")
public class LLMController {

    // 在Controller类顶部添加
    private static final ConcurrentHashMap<String, SseEmitter> taskEmitterMap = new ConcurrentHashMap<>();
    @Resource
    private KanqAiFileService kanqAiFileService;
    @Value("${rerank.context-max-length}")
    private int contextMaxLength;
    @Resource
    private QdrantService qdrantService;
    @Resource
    private LLMService llmService;
    @Resource
    private KanqAiFileChunkService kanqAiFileChunkService;
    @Value("${chunk.top-k}")
    private int topK;

    /**
     * 文件上传接口
     *
     * @param file 待上传的文件
     * @return 返回一个SseEmitter对象，用于服务端发送事件给客户端
     * <p>
     * 该方法的主要功能是接收客户端上传的文件，并通过SseEmitter向客户端发送文件上传的进度事件
     * 使用单线程执行器异步处理文件上传任务，首先发送一个表示文件读取开始的事件，然后读取文件，
     * 文件读取完成后发送一个表示文件读取完成的事件，并关闭SseEmitter连接
     */
    @PostMapping("/upload")
    public SseEmitter upload(@RequestParam("file") MultipartFile file) {
        SseEmitter emitter = new SseEmitter(3 * 60 * 1000L);
        Executors.newSingleThreadExecutor().submit(() -> {
            UploadFileResponse response = new UploadFileResponse();
            response.setUuid(UUID.randomUUID().toString());
            response.setStage("正在读取文件...");
            SseUtil.sseSend(emitter, response);
            int i = llmService.readFile(file);
            if (i == -2) {
                response.setDone(true);
                response.setStage("文件读取失败，文件类型未找到...");
                SseUtil.sseSend(emitter, response);
                emitter.complete();
                return;
            } else if (i == -1) {
                response.setDone(true);
                response.setStage("文件读取失败，请检查文件格式...");
                SseUtil.sseSend(emitter, response);
                emitter.complete();
                return;
            }
            response.setStage("文件读取完成...");
            SseUtil.sseSend(emitter, response);
            emitter.complete();
        });

        return emitter;
    }

    /**
     * 搜索接口
     *
     * @param content 搜索内容
     * @return 返回一个Result对象，包含搜索结果或错误信息
     * <p>
     * 该方法的主要功能是根据用户输入的内容，通过qdrantService获取对应的向量，并搜索相似的向量，
     * 最后返回搜索到的chunkIds列表如果发生异常，则返回错误信息
     */
    @PostMapping("/search")
    public R search(@Param("content") String content) {
        try {
            float[] vector = qdrantService.getEmbedding(content);
            List<Points.ScoredPoint> result = qdrantService.searchVector(vector, 3);
            List<Long> chunkIds = result.stream().map(x -> x.getPayloadMap().get("chunk_id").getIntegerValue()).toList();

            return R.data(chunkIds);
        } catch (Exception e) {
            return R.fail(e.getMessage());
        }
    }


    /**
     * 处理直接聊天请求
     * 该方法接收一个聊天请求，然后通过SSE（Server-Sent Events）方式发送聊天响应
     * 它会根据请求的问题，搜索相关的文档片段，并使用预设的模板和LLM（Large Language Model）服务进行回答
     *
     * @param request 聊天请求对象，包含用户的问题
     * @return 返回一个SseEmitter对象，用于服务器向客户端发送事件
     */
    @PostMapping("/chat")
    public SseEmitter chatDirect(@RequestBody ChatRequest request) {
        return llmService.sseChat(request, taskEmitterMap);
    }


    /**
     * 停止聊天
     *
     * @param taskId
     * @return
     */
    @PostMapping("/stop/{taskId}")
    public R stopChat(@PathVariable String taskId) {
        llmService.stopChat(taskId, taskEmitterMap);
        return R.success("停止成功");
    }


    /**
     * 召回测试
     */
    @GetMapping("/recallTesting")
    public R recallTesting(@RequestParam String question) {
        // 对问题进行分词，获取关键字
        List<String> keywords = HanlpUtils.hanLPSegment(question);
        // 根据关键字在知识库中全文搜索，获取相关文件块ID
        List<Long> chunkIds = kanqAiFileChunkService.fullTextSearch(keywords).stream().map(KanqAiFileChunk::getId).toList();
        //判断chunkIds 是否为空 为空返回
        if (chunkIds.isEmpty()) {
            return R.data(null);
        }
        List<Points.ScoredPoint> result2 = qdrantService.searchVector(qdrantService.getEmbedding(question), topK, chunkIds);
        List<Long> chunkIds2 = result2.stream().map(x -> x.getId().getNum()).toList();
        // 根据文件块ID获取文件及其内容片段
        List<KanqAiFile> fileWithChunks = kanqAiFileService.getFileWithChunksV2(chunkIds2);
        return R.data(fileWithChunks);
    }
}
